use std::cmp::Ordering;
use std::mem;

/// Enum representing the color of a Red-Black Tree node.
#[derive(Clone, Copy, PartialEq, Debug)]
enum RBTreeNodeColor {
    Red,
    Black,
}

/// Enum representing the side of a Red-Black Tree node.
#[derive(Clone, Copy, PartialEq, Debug)]
enum RBTreeNodeSide {
    Left,
    Right,
}

/// Struct representing a node in the Red-Black Tree.
#[derive(Debug)]
struct RBTreeNode<K, V> {
    color: RBTreeNodeColor,
    key: K,
    value: V,
    parent: Option<*mut RBTreeNode<K, V>>,
    children: [Option<*mut RBTreeNode<K, V>>; 2],
}

/// Struct representing the Red-Black Tree.
#[derive(Debug)]
pub struct RBTree<K, V> {
    root_node: Option<*mut RBTreeNode<K, V>>,
    compare_func: fn(&K, &K) -> Ordering,
    num_nodes: usize,
}

impl<K, V> RBTree<K, V> {
    /// Creates a new Red-Black Tree with the given comparison function.
    pub fn new(compare_func: fn(&K, &K) -> Ordering) -> Self {
        RBTree {
            root_node: None,
            compare_func,
            num_nodes: 0,
        }
    }

    /// Frees the entire tree, including all nodes.
    pub fn free(&mut self) {
        if let Some(root) = self.root_node {
            self.free_subtree(root);
        }
        self.root_node = None;
    }

    /// Recursively frees the subtree rooted at the given node.
    fn free_subtree(&mut self, node: *mut RBTreeNode<K, V>) {
        unsafe {
            if let Some(left) = (*node).children[RBTreeNodeSide::Left as usize] {
                self.free_subtree(left);
            }
            if let Some(right) = (*node).children[RBTreeNodeSide::Right as usize] {
                self.free_subtree(right);
            }
            Box::from_raw(node);
        }
    }

    /// Inserts a new key-value pair into the tree.
    pub fn insert(&mut self, key: K, value: V) {
        let node = Box::into_raw(Box::new(RBTreeNode {
            color: RBTreeNodeColor::Red,
            key,
            value,
            parent: None,
            children: [None, None],
        }));

        let mut parent = None;
        let mut rover = self.root_node;

        while let Some(current) = rover {
            parent = Some(current);
            let side = match (self.compare_func)(&key, unsafe { &(*current).key }) {
                Ordering::Less => RBTreeNodeSide::Left,
                _ => RBTreeNodeSide::Right,
            };
            rover = unsafe { (*current).children[side as usize] };
        }

        unsafe {
            (*node).parent = parent;
        }

        if let Some(parent) = parent {
            let side = match (self.compare_func)(&key, unsafe { &(*parent).key }) {
                Ordering::Less => RBTreeNodeSide::Left,
                _ => RBTreeNodeSide::Right,
            };
            unsafe {
                (*parent).children[side as usize] = Some(node);
            }
        } else {
            self.root_node = Some(node);
        }

        self.insert_case1(node);
        self.num_nodes += 1;
    }

    /// Returns the number of nodes in the tree.
    pub fn num_entries(&self) -> usize {
        self.num_nodes
    }

    /// Returns the root node of the tree.
    pub fn root_node(&self) -> Option<*mut RBTreeNode<K, V>> {
        self.root_node
    }

    /// Returns the key of the given node.
    pub fn node_key(&self, node: *mut RBTreeNode<K, V>) -> &K {
        unsafe { &(*node).key }
    }

    /// Returns the value of the given node.
    pub fn node_value(&self, node: *mut RBTreeNode<K, V>) -> &V {
        unsafe { &(*node).value }
    }

    /// Returns the child of the given node on the specified side.
    pub fn node_child(&self, node: *mut RBTreeNode<K, V>, side: RBTreeNodeSide) -> Option<*mut RBTreeNode<K, V>> {
        unsafe { (*node).children[side as usize] }
    }

    /// Returns the parent of the given node.
    pub fn node_parent(&self, node: *mut RBTreeNode<K, V>) -> Option<*mut RBTreeNode<K, V>> {
        unsafe { (*node).parent }
    }

    /// Returns the sibling of the given node.
    fn node_sibling(&self, node: *mut RBTreeNode<K, V>) -> Option<*mut RBTreeNode<K, V>> {
        unsafe {
            let parent = (*node).parent?;
            let side = self.node_side(node);
            (*parent).children[1 - side as usize]
        }
    }

    /// Returns the uncle of the given node.
    fn node_uncle(&self, node: *mut RBTreeNode<K, V>) -> Option<*mut RBTreeNode<K, V>> {
        unsafe {
            let parent = (*node).parent?;
            self.node_sibling(parent)
        }
    }

    /// Returns the side of the given node relative to its parent.
    fn node_side(&self, node: *mut RBTreeNode<K, V>) -> RBTreeNodeSide {
        unsafe {
            let parent = (*node).parent.expect("Node has no parent");
            if (*parent).children[RBTreeNodeSide::Left as usize] == Some(node) {
                RBTreeNodeSide::Left
            } else {
                RBTreeNodeSide::Right
            }
        }
    }

    /// Replaces node1 with node2 at its parent.
    fn node_replace(&mut self, node1: *mut RBTreeNode<K, V>, node2: Option<*mut RBTreeNode<K, V>>) {
        if let Some(node2) = node2 {
            unsafe {
                (*node2).parent = (*node1).parent;
            }
        }

        if let Some(parent) = unsafe { (*node1).parent } {
            let side = self.node_side(node1);
            unsafe {
                (*parent).children[side as usize] = node2;
            }
        } else {
            self.root_node = node2;
        }
    }

    /// Rotates the subtree rooted at the given node in the specified direction.
    fn rotate(&mut self, node: *mut RBTreeNode<K, V>, direction: RBTreeNodeSide) {
        let new_root = unsafe { (*node).children[1 - direction as usize].expect("Rotation node has no child") };

        self.node_replace(node, Some(new_root));

        unsafe {
            (*node).children[1 - direction as usize] = (*new_root).children[direction as usize];
            (*new_root).children[direction as usize] = Some(node);
            (*node).parent = Some(new_root);

            if let Some(child) = (*node).children[1 - direction as usize] {
                (*child).parent = Some(node);
            }
        }
    }

    /// Handles the first case of insertion in a Red-Black Tree.
    fn insert_case1(&mut self, node: *mut RBTreeNode<K, V>) {
        if unsafe { (*node).parent.is_none() } {
            unsafe {
                (*node).color = RBTreeNodeColor::Black;
            }
        } else {
            self.insert_case2(node);
        }
    }

    /// Handles the second case of insertion in a Red-Black Tree.
    fn insert_case2(&mut self, node: *mut RBTreeNode<K, V>) {
        if unsafe { (*(*node).parent.unwrap()).color == RBTreeNodeColor::Red } {
            self.insert_case3(node);
        }
    }

    /// Handles the third case of insertion in a Red-Black Tree.
    fn insert_case3(&mut self, node: *mut RBTreeNode<K, V>) {
        let uncle = self.node_uncle(node);
        let grandparent = unsafe { (*(*node).parent.unwrap()).parent.unwrap() };

        if let Some(uncle) = uncle {
            if unsafe { (*uncle).color == RBTreeNodeColor::Red } {
                unsafe {
                    (*(*node).parent.unwrap()).color = RBTreeNodeColor::Black;
                    (*uncle).color = RBTreeNodeColor::Black;
                    (*grandparent).color = RBTreeNodeColor::Red;
                }
                self.insert_case1(grandparent);
            } else {
                self.insert_case4(node);
            }
        } else {
            self.insert_case4(node);
        }
    }

    /// Handles the fourth case of insertion in a Red-Black Tree.
    fn insert_case4(&mut self, node: *mut RBTreeNode<K, V>) {
        let side = self.node_side(node);
        let parent = unsafe { (*node).parent.unwrap() };
        let grandparent = unsafe { (*parent).parent.unwrap() };

        if side != self.node_side(parent) {
            self.rotate(parent, side);
            self.insert_case5(parent);
        } else {
            self.insert_case5(node);
        }
    }

    /// Handles the fifth case of insertion in a Red-Black Tree.
    fn insert_case5(&mut self, node: *mut RBTreeNode<K, V>) {
        let parent = unsafe { (*node).parent.unwrap() };
        let grandparent = unsafe { (*parent).parent.unwrap() };
        let side = self.node_side(node);

        self.rotate(grandparent, 1 - side);
        unsafe {
            (*parent).color = RBTreeNodeColor::Black;
            (*grandparent).color = RBTreeNodeColor::Red;
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_insert_and_lookup() {
        let mut tree = RBTree::new(|a: &i32, b: &i32| a.cmp(b));
        tree.insert(10, "ten");
        tree.insert(20, "twenty");
        tree.insert(30, "thirty");

        assert_eq!(tree.num_entries(), 3);

        let root = tree.root_node().unwrap();
        assert_eq!(tree.node_key(root), &20);
        assert_eq!(tree.node_value(root), &"twenty");

        let left = tree.node_child(root, RBTreeNodeSide::Left).unwrap();
        assert_eq!(tree.node_key(left), &10);
        assert_eq!(tree.node_value(left), &"ten");

        let right = tree.node_child(root, RBTreeNodeSide::Right).unwrap();
        assert_eq!(tree.node_key(right), &30);
        assert_eq!(tree.node_value(right), &"thirty");
    }
}
